package com.study.Utils;

import org.elasticsearch.action.search.SearchRequestBuilder;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.client.Client;
import org.elasticsearch.search.aggregations.Aggregation;
import org.elasticsearch.search.aggregations.AggregationBuilders;
import org.elasticsearch.search.aggregations.bucket.terms.StringTerms;
import org.elasticsearch.search.aggregations.bucket.terms.TermsAggregationBuilder;
import org.springframework.data.elasticsearch.core.ElasticsearchTemplate;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * ES 聚合查询 util
 *
 * @author huang
 */
@Component
public class ESAggregationUtil {

    @Resource
    private ElasticsearchTemplate elasticsearchTemplate;

    /**
     * 分组求和
     *
     * @param indexName indexName
     * @param typeName  type
     * @param fields    分组字段
     * @param size      分组数, 默认仅返回最多 10 个组
     */
    public HashMap<String, Long> groupBy(String indexName, String typeName, List<String> fields, Integer size) {
        HashMap<String, Long> resultMap = new HashMap<>(1);
        // 设置 index, type
        SearchRequestBuilder searchRequestBuilder = elasticsearchTemplate.getClient().prepareSearch(indexName).setTypes(typeName);
        // 构建分组, 指定 size
        TermsAggregationBuilder group = AggregationBuilders.terms("group").size(size);
        // 添加分组字段
        fields.forEach(group::field);
        searchRequestBuilder.addAggregation(group);
        // 执行查询
        SearchResponse searchResponse = searchRequestBuilder.execute().actionGet();
        // 查询结果转 map
        Map<String, Aggregation> map = searchResponse.getAggregations().asMap();
        fields.forEach(item -> {
            StringTerms count = (StringTerms) map.get("group");
            // 有多少个组
            List<StringTerms.Bucket> buckets = count.getBuckets();
            for (StringTerms.Bucket bucket : buckets) {
                // key, 字段名
                String keyAsString = bucket.getKeyAsString();
                // count
                long docCount = bucket.getDocCount();

                resultMap.put(keyAsString, docCount);
            }
        });
        return resultMap;
    }
}
